package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

const TOTAL = 10000

var (
	//需要测试的url
	abURL string
	//需要发送的post 文件路径
	postFile string
	//total request
	totalReq int
	//需要的auth
	auth    string
	failErr = errors.New("request fail")
)

func init() {
	flag.StringVar(&abURL, "url", "", "api to test with http")
	flag.StringVar(&postFile, "path", "", "post json file to ab")
	flag.StringVar(&auth, "auth", "", "required authorization")
	flag.IntVar(&totalReq, "n", TOTAL, "total request")
}

func main() {
	flag.Parse()
	if abURL == "" {
		flag.Usage()
		os.Exit(1)
	}
	var mdText []string
	for i := 10; i <= 1000; i <<= 1 {
		api := makeCMD(postFile, i, totalReq, abURL, auth)
		cmd := exec.Command("bash", "-c", api)
		println("ab test cmd: ", api)
		//创建获取命令输出管道
		stdout, err := cmd.StdoutPipe()
		if err != nil {
			fmt.Printf("Error:can not obtain stdout pipe for command:%s\n", err)
			return
		}

		//执行命令
		if err := cmd.Start(); err != nil {
			fmt.Println("Error:The command is err,", err)
			return
		}

		//读取所有输出
		text, failed := handleText(stdout)
		if err := cmd.Wait(); err != nil {
			fmt.Println("wait:", err.Error())
			//退出前打印当前的统计
			printQPS(mdText)
			return
		}

		if text != nil {
			md := fmt.Sprintf("%d|%d|%s", i, totalReq, text["Requests per second"])
			mdText = append(mdText, md)
		}
		if failed == failErr {
			println("max qps with concurrency ", i)
			break
		}
	}
	printQPS(mdText)
}

func printQPS(mdText []string) {
	if mdText != nil {
		fullText := "并发数|请求总数|qps\n-|-|-\n" + strings.Join(mdText, "\n")
		println(fullText)
		return
	}
	println("no qps collected")
}

//生成 ab 命令
func makeCMD(postFile string, concurrency int, totalRequest int, url string, auth string) string {
	if len(url) < 4 || strings.ToLower(url[:4]) != "http" {
		url = "https://" + url
	}
	if auth != "" {
		//添加授权
		auth = "-H \"Authorization: " + auth + "\""
	}
	if postFile == "" {
		return fmt.Sprintf("ab -c %d -n %d %s \"%s\"", concurrency, totalRequest, auth, url)
	}
	p, err := filepath.Abs(postFile)
	if err != nil {
		panic("invalid path:" + err.Error())
	}
	return fmt.Sprintf("ab -p %s -T application/json -c %d -n %d %s \"%s\"", p, concurrency, totalRequest, auth, url)
}

// read print body
func handleText(stdout io.ReadCloser) (kv map[string]string, err error) {
	out := bufio.NewReader(stdout)
	kv = make(map[string]string)
	for {
		bytes, _, err := out.ReadLine()
		if err != nil {
			if err.Error() == "EOF" {
				return kv, nil
			}
			//err in read
			log.Print(err)
			return nil, err
		}
		line := string(bytes)
		if len(line) == 0 || !strings.Contains(line, ":") {
			continue
		}

		temp := strings.Split(line, ":")
		var key, value string
		for k, v := range temp {
			if k == 0 {
				key = v
			} else if k == 1 {
				value = v
			}
		}

		// Failed requests
		value = strings.ReplaceAll(value, " ", "")
		if key == "Failed requests" && value != "0" {
			//请求失败
			return nil, failErr
		}
		kv[key] = value
	}

}
